#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
动量轮动策略回测系统 (L21H5版本)

策略说明：
1. 动量计算周期：21个交易日（约1个月）
   - 选择21天作为动量计算周期，可以较好地捕捉中期趋势
   - 避免过短周期带来的噪音和过度交易
   - 避免过长周期导致信号滞后

2. 持仓周期：5个交易日（1周）
   - 每周调仓一次，平衡交易成本和策略灵敏度
   - 降低换手率，减少交易成本
   - 避免频繁交易带来的滑点损失

3. 策略特点：
   - 基于动量的趋势跟踪
   - 多资产轮动配置
   - 自动化的交易执行
   - 风险分散化投资

4. 参数说明：
   - LOOKBACK: 21 (动量计算周期)
   - HOLD_PERIOD: 5 (持仓周期)
   - ETF_LIST: 包含多个行业和主题ETF

5. 使用方法：
   python momentum_rotation_backtesting_l21h5.py --start_date 20200101 --end_date 20241231
   python momentum_rotation_backtesting_l21h5.py --daily_signal  # 获取当日交易信号
"""

import akshare as ak
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import argparse
from datetime import datetime, timedelta

# 设置中文字体（任选其一）
plt.rcParams['font.sans-serif'] = ['SimHei']  # Windows系统黑体
# plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']  # macOS
# plt.rcParams['font.sans-serif'] = ['WenQuanYi Zen Hei']  # Linux

# 解决负号显示问题
plt.rcParams['axes.unicode_minus'] = False

# 参数设置
ETF_LIST = ['899050', '588000']
ETF_NAMES = {
    '899050': '北证50ETF',
    '588000': '科创50ETF'
}

LOOKBACK = 21  # 动量计算周期（21个交易日≈1个月）
HOLD_PERIOD = 5  # 持仓周期（交易日）

def get_etf_data(symbol, start_date, end_date):
    """获取ETF历史数据"""
    df = ak.fund_etf_hist_em(symbol=symbol, period="daily",
                             start_date=start_date, end_date=end_date,
                             adjust="hfq")  # 后复权
    df['date'] = pd.to_datetime(df['日期'])
    df.set_index('date', inplace=True)
    return df['收盘'].rename(symbol)

def calculate_momentum(df, window=LOOKBACK):
    """计算动量信号"""
    return df.pct_change(window).shift(1)  # 避免未来函数

def generate_signal(momentum_df):
    """生成交易信号"""
    rank = momentum_df.rank(axis=1, ascending=False)
    return rank <= 1  # 选择动量第1名

def get_etf_display_name(etf_code):
    """获取ETF的显示名称（代码-名称）"""
    etf_name = ETF_NAMES.get(etf_code, etf_code)
    return f"{etf_code}-{etf_name}"

def run_momentum_rotation(start_date, end_date):
    """运行动量轮动策略回测"""
    # 获取数据
    data = pd.concat([get_etf_data(etf, start_date, end_date) for etf in ETF_LIST], axis=1).dropna()
    
    # 计算动量信号
    momentum = calculate_momentum(data)
    
    # 生成交易信号
    signal = generate_signal(momentum)
    
    # 执行回测
    portfolio = pd.DataFrame(index=data.index)
    portfolio['position'] = signal.idxmax(axis=1).shift(1)  # T+1日执行
    portfolio['return'] = data.apply(lambda x: x.pct_change(), axis=0) \
        .where(signal.shift(1)).sum(axis=1)
    portfolio['cum_return'] = (1 + portfolio['return']).cumprod()
    
    # 打印交易记录
    print_trade_records(portfolio, data)
    
    # 可视化
    plt.figure(figsize=(12, 6))
    portfolio['cum_return'].plot(label='策略净值')
    (data / data.iloc[0]).mean(axis=1).plot(label='等权基准', linestyle='--')
    plt.title(f'动量轮动策略表现 (回测周期: {start_date}至{end_date})')
    plt.legend()
    plt.show()
    
    # 绩效统计
    annual_return = portfolio['return'].mean() * 252
    max_drawdown = (portfolio['cum_return'] / portfolio['cum_return'].cummax() - 1).min()
    print(f"年化收益: {annual_return:.2%} | 最大回撤: {max_drawdown:.2%}")
    
    return portfolio, data, signal


def print_trade_records(portfolio, data):
    """打印交易记录（最终修正版）"""
    # 获取持仓变化点
    position_changes = portfolio['position'].ne(portfolio['position'].shift()).astype(int)
    rebalance_dates = position_changes[position_changes == 1].index

    print("\n===== 交易记录（最终修正版） =====")
    print("序号 | 买入日期 | 买入ETF | 卖出日期 | 卖出ETF | 持有天数 | 收益率")
    print("-" * 100)

    # 初始化首个持仓
    if len(rebalance_dates) > 0:
        first_date = rebalance_dates[0]
        first_etf = portfolio.loc[first_date, 'position']
        print(
            f"  1 | {first_date.strftime('%Y-%m-%d')} | {get_etf_display_name(first_etf)} | {'初始买入':^10} | {'-':^8} | {'-':^8} | {'-':^8}")

    # 遍历调仓点生成记录
    trade_count = 1
    for i in range(1, len(rebalance_dates)):
        # 获取交易时间区间
        buy_date = rebalance_dates[i - 1]  # 本次买入日期（即上次卖出日期）
        sell_date = rebalance_dates[i]  # 本次卖出日期

        # 获取交易标的
        sold_etf = portfolio.loc[buy_date, 'position']  # 卖出的是上次买入的ETF
        bought_etf = portfolio.loc[sell_date, 'position']  # 本次买入的ETF

        # 计算持有信息
        holding_days = (sell_date - buy_date).days
        buy_price = data.loc[buy_date, sold_etf]
        sell_price = data.loc[sell_date, sold_etf]
        return_rate = (sell_price / buy_price - 1) * 100

        # 打印交易记录
        trade_count += 1
        print(f"{trade_count:3d} | {buy_date.strftime('%Y-%m-%d')} | {get_etf_display_name(sold_etf)} | "
              f"{sell_date.strftime('%Y-%m-%d')} | {get_etf_display_name(bought_etf)} | "
              f"{holding_days:^8d} | {return_rate:>6.2f}%")

    # 处理最后一个持仓
    if len(rebalance_dates) > 0:
        last_buy_date = rebalance_dates[-1]
        last_etf = portfolio.loc[last_buy_date, 'position']
        holding_days = (data.index[-1] - last_buy_date).days

        buy_price = data.loc[last_buy_date, last_etf]
        sell_price = data.loc[data.index[-1], last_etf]
        return_rate = (sell_price / buy_price - 1) * 100
        
        print(f"{trade_count + 1:3d} | {last_buy_date.strftime('%Y-%m-%d')} | {get_etf_display_name(last_etf)} | "
              f"{data.index[-1].strftime('%Y-%m-%d')} | {'持有中':^8} | {holding_days:^8d} | {return_rate:>6.2f}%")

    print("-" * 100)
    print(f"总交易次数: {trade_count}")

def get_daily_signal():
    """获取当日交易信号"""
    # 计算日期范围：获取足够的历史数据来计算动量
    end_date = datetime.now().strftime('%Y%m%d')
    # 获取LOOKBACK+30天的数据，确保有足够的历史数据计算动量
    start_date = (datetime.now() - timedelta(days=LOOKBACK*2)).strftime('%Y%m%d')
    
    print(f"\n===== 动量轮动策略 - 每日信号 ({datetime.now().strftime('%Y-%m-%d')}) =====")
    print(f"数据获取范围: {start_date} 至 {end_date}")
    
    try:
        # 获取数据
        data = pd.concat([get_etf_data(etf, start_date, end_date) for etf in ETF_LIST], axis=1).dropna()
        
        # 计算动量信号
        momentum = calculate_momentum(data)
        
        # 生成交易信号
        signal = generate_signal(momentum)
        
        # 获取最新信号
        latest_date = data.index[-1]
        latest_momentum = momentum.loc[latest_date]
        latest_signal = signal.loc[latest_date]
        
        # 打印动量排名
        print("\n动量排名:")
        momentum_rank = latest_momentum.rank(ascending=False)
        for etf in ETF_LIST:
            rank = momentum_rank[etf]
            momentum_value = latest_momentum[etf] * 100
            etf_display = get_etf_display_name(etf)
            print(f"{etf_display}: 排名 {rank:.0f}, 动量 {momentum_value:.2f}%")
        
        # 打印当前持仓建议
        current_position = signal.idxmax(axis=1).shift(1).iloc[-1]
        current_position_display = get_etf_display_name(current_position)
        print(f"\n当前持仓建议: {current_position_display}")
        
        # 计算距离下次调仓的天数
        position_changes = signal.idxmax(axis=1).ne(signal.idxmax(axis=1).shift()).astype(int)
        last_rebalance = position_changes[position_changes == 1].index[-1]
        days_since_rebalance = (latest_date - last_rebalance).days
        
        print(f"距离上次调仓: {days_since_rebalance} 天")
        print(f"距离下次调仓: {HOLD_PERIOD - days_since_rebalance} 天")
        
        # 打印历史表现
        portfolio = pd.DataFrame(index=data.index)
        portfolio['position'] = signal.idxmax(axis=1).shift(1)
        portfolio['return'] = data.apply(lambda x: x.pct_change(), axis=0).where(signal.shift(1)).sum(axis=1)
        portfolio['cum_return'] = (1 + portfolio['return']).cumprod()
        
        # 计算最近表现
        recent_return = (portfolio['cum_return'].iloc[-1] / portfolio['cum_return'].iloc[-20] - 1) * 100 if len(portfolio) > 20 else 0
        annual_return = portfolio['return'].mean() * 252 * 100
        
        print(f"\n策略表现:")
        print(f"最近20日收益率: {recent_return:.2f}%")
        print(f"年化收益率: {annual_return:.2f}%")
        
    except Exception as e:
        print(f"获取信号时出错: {e}")

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='动量轮动策略回测')
    parser.add_argument('--start_date', type=str, default='20241001',
                      help='回测开始日期 (YYYYMMDD)')
    parser.add_argument('--end_date', type=str, 
                      default=datetime.now().strftime('%Y%m%d'),
                      help='回测结束日期 (YYYYMMDD)')
    parser.add_argument('--daily_signal', action='store_true',
                      help='获取当日交易信号')
    return parser.parse_args()

def main():
    """主函数"""
    args = parse_args()
    
    if args.daily_signal:
        get_daily_signal()
    else:
        run_momentum_rotation(args.start_date, args.end_date)

if __name__ == "__main__":
    main()
